import json
from typing import Tuple

from openai import OpenAI

client = OpenAI(
    api_key='sk-wUms3yGLjiImG1TI967c9658Ca6b4cBa9583B7Ae942305C8',
    base_url='https://api.ai-gaochao.cn/v1'
)


class Chat:
    def __init__(self):
        self.count, self.dict_, self.sentence_list = 1, {}, []
        with open('record.json', 'a+', encoding='utf-8') as f:
            if f.read():
                self.dict_ = json.load(f)
                self.count = max(self.dict_.keys()) + 1
        self.sentence_list.append({'role': 'system', 'content': ''})

    def ask(self, prompt):
        self.sentence_list.append({"role": "user", "content": prompt})
        response = client.chat.completions.create(model="gpt-3.5-turbo", messages=self.sentence_list)
        answer = json.loads(response.model_dump_json())['choices'][0]['message']['content']
        self.sentence_list.append({"role": "assistant", "content": answer})
        self.dict_[self.count] = self.sentence_list
        with open("record.json", "w", encoding='utf-8') as f:
            json.dump(self.dict_, f, ensure_ascii=False, indent=2)
        return answer


def generate_news_main(prompt_file_path: str, key_info: str) -> Tuple[str, str]:
    model_GPT = 'gpt-3.5-turbo-16k'

    # model_GPT = "gpt-4-1106-preview"

    # 读取系统消息
    with open(prompt_file_path, 'r', encoding='utf-8') as file:  # 这是预设的通用Prompt信息
        system_message = file.read()

    # 消息列表，包含系统消息和用户消息
    messages = [
        {
            "role": "system",
            "content": system_message
        },
        {
            "role": "user",
            "content": key_info
        }
    ]

    # 调用OpenAI的API生成回复
    response = client.chat.completions.create(
        messages=messages,
        model=model_GPT,
    )

    # 解析响应
    answer = json.loads(response.model_dump_json())['choices'][0]['message']['content']

    # 假设生成的答案的第一行是标题，剩下的是正文
    title, _, body = answer.partition('\n')

    # 打印结果
    print("新闻标题:", title)
    print("新闻正文:", body)

    return title, body
